import numpy as np
import supervision as sv
from ultralytics import YOLO

from trackers import DeepSORTTracker, DeepSORTFeatureExtractor

model = YOLO("yolov8l", task="detect", verbose=False)
feature_extractor = DeepSORTFeatureExtractor.from_timm(
    model_name="mobilenetv4_conv_small.e1200_r224_in1k", pretrained=True
)
tracker = DeepSORTTracker(feature_extractor=feature_extractor)
box_annotator = sv.BoxAnnotator()
label_annotator = sv.LabelAnnotator()


def callback(frame: np.ndarray, _: int):
    results = model(frame, verbose=False, imgsz=1280)[0]
    detections = sv.Detections.from_ultralytics(results)
    detections = tracker.update(detections, frame)

    labels = [
        f"#{tracker_id} {results.names[class_id]} {confidence:.2f}"
        for tracker_id, class_id, confidence in zip(
            detections.tracker_id, detections.class_id, detections.confidence
        )
    ]
    annotated_image = frame.copy()
    annotated_image = box_annotator.annotate(annotated_image, detections)
    annotated_image = label_annotator.annotate(annotated_image, detections, labels)

    return annotated_image


sv.process_video(
    source_path="data/traffic_video.mp4",
    target_path="data/out.mp4",
    callback=callback,
    max_frames=100,
)